# -*- coding:utf-8 -*-
# !/bin/env python3

import os
import sys
import struct
import pickle
import soundfile as sf
import numpy as np
import argparse


class DataUnPack(object):
    def __init__(self):
        # paras
        self._head_bytes_size = 24
        self._data_info_map = {
            "check_status": 0,
            "date_type": {
                0: "audio_pickle",
                1: "seq_path_pickle",
                2: "no_define",
            },
            "pack_version": 1.0,
            "data_length": 0,
        }

        self.update_desc_base_path = "@DESC_BASE_PATH@"
        self.aim_desc_base_path = "/apdcephfs_cq2/share_1297902/speech_data/pack_data_multilingual"
        self.update_audio_base_path = "@AUDIO_BASE_PATH@"
        self.aim_audio_base_path = "/apdcephfs_cq2/share_1297902/speech_data/pack_data_seq"

        self.fp = None
        self.fp_file = None

    def read_head_fp(self, fp):
        tmp_data = fp.read(self._head_bytes_size)
        info_list = struct.unpack("iifq", tmp_data)
        if info_list[0] != self._data_info_map["check_status"]:
            raise Exception("wrong read_head with check_status:%s" \
                            % (info_list[0]))

        if info_list[1] not in self._data_info_map["date_type"]:
            raise Exception("wrong read_head with date_type: %s" \
                            % (info_list[1]))

        if info_list[2] > self._data_info_map["pack_version"]:
            raise Exception("wrong read_head with pack_version: %s" \
                            % (info_list[2]))

        if info_list[3] <= self._data_info_map["data_length"]:
            raise Exception("wrong read_head with data_length: %s" \
                            % (info_list[3]))
        return fp, {"date_type": info_list[1], \
                    "pack_version": info_list[2], "data_length": info_list[3]}

    def del_feff(self, input_string):
        return input_string.encode('utf-8').decode('utf-8-sig')

    def load_pcm_buffer(self, pcm_buffer):
        audio = np.frombuffer(pcm_buffer, dtype='h')
        audio = audio.astype('float32')
        audio = audio / 32767.0
        return audio

    def read_buffer_from_seq(self, input_str):
        tmp_info = input_str.strip().split("|")
        if len(tmp_info) < 2:
            return -1, None

        if "seq" == tmp_info[0]:
            file_path = tmp_info[1]
            file_seek = int(tmp_info[2])
            file_length = int(tmp_info[3])
            fp = open(file_path, 'rb')
            fp.seek(file_seek)
            data = fp.read(file_length)
            fp.close()
            if len(data) != file_length:
                return -2, None

            data = self.load_pcm_buffer(data)
        else:
            return -3, None
        return 0, data
    def read_buffer_from_seq_xlgeng(self, file_path):
        fp = open(file_path, 'rb')
        data = fp.read()
        fp.close()
        data = self.load_pcm_buffer(data)
        return 0, data

    def save_audio(self, audio_buffer, path, aim_type="pcm", sample_rate=16000):
        if "pcm" == aim_type:
            data = np.minimum(np.maximum(audio_buffer * 32767, \
                                         -32767 * np.ones(audio_buffer.shape)), 32767 * np.ones(audio_buffer.shape))
            data = data.astype('h').tostring()
            open(path, 'wb').write(data)
        elif "wav" == aim_type:
            sf.write(path, audio_buffer, sample_rate)
        else:
            return -1
        return 0

    def read_desc_line(self, input_line):
        input_line = self.del_feff(input_line).strip()
        tmp_info = input_line.split()
        if len(tmp_info) < 2:
            return {}

        aim_file = tmp_info[1]
        fp_begin = int(tmp_info[2])

        if self.fp is not None:
            if aim_file != self.fp_file:
                self.fp.close()

                self.fp_file = aim_file
                self.fp = open(aim_file, 'rb')
        else:
            self.fp_file = aim_file
            self.fp = open(aim_file, 'rb')

        self.fp.seek(fp_begin)
        self.fp, info_dict = self.read_head_fp(self.fp)
        data_dict = pickle.loads(self.fp.read(info_dict["data_length"]))
        return data_dict

    def read_desc_line_independence(self, input_line):
        input_line = self.del_feff(input_line).strip()
        tmp_info = input_line.split()
        if len(tmp_info) < 2:
            return {}

        aim_file = tmp_info[1]
        fp_begin = int(tmp_info[2])
        aim_file = aim_file.replace("apdcephfs_cq8/share_2906397/users", "apdcephfs_sh7/share_6093881/hunyuan")
        fp = open(aim_file, 'rb')
        fp.seek(fp_begin)
        fp, info_dict = self.read_head_fp(fp)
        data_dict = pickle.loads(fp.read(info_dict["data_length"]))
        fp.close()
        return data_dict

    def close_fp(self):
        if self.fp is not None:
            self.fp.close()
        return 0

    def read_desc_line_with_show(self, input_line):
        # import json
        # input_line = json.load(input_line)
        input_line = self.del_feff(input_line).strip()
        # print(input_line)
        tmp_info = input_line.split()
        # print(tmp_info)
        if len(tmp_info) < 2:
            return -1
        # print(tmp_info[2])
        aim_file = tmp_info[1]
        fp_begin = int(tmp_info[2])
        aim_file = aim_file.replace("apdcephfs_cq8/share_2906397/users", "apdcephfs_sh7/share_6093881/hunyuan")

        fp = open(aim_file, 'rb')
        fp.seek(fp_begin)
        fp, info_dict = self.read_head_fp(fp)

        if 1 == info_dict["date_type"]:
            data_dict = pickle.loads(fp.read(info_dict["data_length"]))

            for lang in data_dict["lang_type"]:
                for item in data_dict[lang]:
                    if item in ["pcm16k16bit", "pcm24k16bit"]:
                        for i in range(len(data_dict[lang][item])):
                            seq_info = data_dict[lang][item][i][0]
                            seq_info = seq_info.replace(self.update_audio_base_path, self.aim_audio_base_path)
                            data_dict[lang][item][i][0] = seq_info
            print(data_dict)
        fp.close()
        return 0

    def read_desc_file(self, input_desc_file):
        # import pdb;pdb.set_trace()
        fp_input = open(input_desc_file, 'r')
        for i in range(1):
            input_line = fp_input.readline()
            if not input_line:
                fp_input.close()
                break
            """
            input_line = input_line.replace(self.update_desc_base_path, self.aim_desc_base_path)
            #self.read_desc_line_with_show(input_line)
            desc_data_dict = self.read_desc_line_independence(input_line)
            #print(data_dict)

            lang = desc_data_dict["lang_type"]
            print(lang)
            audio_path = desc_data_dict[lang[0]]["pcm16k16bit"][0][0]
            audio_path = audio_path.replace(
                self.desc_data_unpack.update_audio_base_path,
                self.desc_data_unpack.aim_audio_base_path,
            )
            print(audio_path)
            _, audio =  self.desc_data_unpack.read_buffer_from_seq(audio_path)
            print(len(audio))


            # for save audio
            data_dict = self.read_desc_line(input_line)
            print(data_dict["key_original"])
            for lang in data_dict["lang_type"]:
                for item in data_dict[lang]:
                    if item in ["text"]:
                        text = data_dict[lang][item][0]
                        print(lang, text)
                    elif item in ["pcm16k16bit"]:
                        seq_info = data_dict[lang][item][0][0]
                        seq_info = seq_info.replace(self.update_audio_base_path, self.aim_audio_base_path)
                        rtn, audio_data = self.read_buffer_from_seq(seq_info)
                        if rtn == 0:
                            self.save_audio(audio_data, "./%s.wav" % data_dict["key_original"], "wav")
                        print(lang, seq_info)
                    else:
                        continue
            """

        return 0


if __name__ == "__main__":
    # Argparse
    parser = argparse.ArgumentParser(description='DataUnPack')
    # Read yaml
    parser.add_argument("--desc", required=True,
                        default="/dockerdata/rainiejjli/output_zh_cszhen_en_train.desc.shuf_00_nobiaoji",
                        help="input desc file")

    args = parser.parse_args()
    desc_data_unpack = DataUnPack()
    # task.read_desc_file(args.desc)
    fp_input = open(args.desc, 'r')
    for i in range(1):
        input_line = fp_input.readline()
        if not input_line:
            fp_input.close()
            break

        desc_data_dict = desc_data_unpack.read_desc_line_independence(input_line)
        # print()
        lang = desc_data_dict["lang_type"]
        asr_text = desc_data_dict[lang[0]]["text"][0]
        print("source lang: ", lang[0])
        print("source text: ", asr_text)
        audio_path = desc_data_dict[lang[0]]["pcm16k16bit"][0][0]
        audio_path = audio_path.replace(
            desc_data_unpack.update_audio_base_path,
            desc_data_unpack.aim_audio_base_path,
        )
        print(audio_path)
        audio_path = "seq|/apdcephfs_cq8/share_2906397/users/rainiejjli/datasets/zh/out_zh_part1_scp27/pack_file_name.7.seq|8046904678|66448"
        audio_path = audio_path.replace("apdcephfs_cq8/share_2906397/users", "apdcephfs_sh7/share_6093881/hunyuan")
        _, audio = desc_data_unpack.read_buffer_from_seq(audio_path)
        print(audio)
        sf.write("tmp.wav", audio, 16000)
        print(len(audio))